"""
========================
画图核心模块
------------------------
Author: Lilidream
Date: 2025-03-15
========================
"""
from dataclasses import dataclass, asdict
from matplotlib import pyplot as plt
from matplotlib.axes import Axes
from matplotlib.transforms import Bbox, BboxTransformFrom
from numpy import ndarray
from matplotlib.figure import Figure
from typing import Union, Literal
from abc import ABC, abstractmethod

AxesGroup = Union[Axes, list[Axes], ndarray[Axes, Axes]]

def axes_to_list(axes: AxesGroup) -> list[Axes]:
    """
    将一个或一组 axes 转换为 list[Axes]
    """
    if isinstance(axes, list):
        return axes
    elif isinstance(axes, ndarray):
        return axes.flatten().tolist()
    else:
        return [axes]


def bbox_xy(axes: AxesGroup, fig: Figure=None, 
            tight_box: bool = False) -> Bbox:
    """
    计算一组 axes 最外围的 bounding box, 返回的 bbox 是以 xy 为坐标系的（即0-1）
    """
    axes = axes_to_list(axes)

    left = 1
    bottom = 1
    right = 0
    top = 0

    if tight_box:
        fig_width_px, fig_height_px = fig.get_size_inches() * fig.dpi
        fig_trans = BboxTransformFrom(Bbox.from_extents(0, 0, fig_width_px, fig_height_px))

    for ax in axes:
        if bool(tight_box):
            bbox = ax.get_tightbbox(fig.canvas.get_renderer())
            lower_left = fig_trans.transform_point((bbox.x0, bbox.y0))
            upper_right = fig_trans.transform_point((bbox.x1, bbox.y1))
        else:
            bbox = ax.get_position()
            lower_left = (bbox.x0, bbox.y0)
            upper_right = (bbox.x1, bbox.y1)

        left = min(left, lower_left[0])
        bottom = min(bottom, lower_left[1])
        right = max(right, upper_right[0])
        top = max(top, upper_right[1])

    return Bbox.from_extents(left, bottom, right, top)


def shift_axes(axes: AxesGroup, dx: float = 0, dy: float = 0):
    """
    移动一组 axes
    """
    axes = axes_to_list(axes)
    for ax in axes:
        bbox = ax.get_position()
        ax.set_position([bbox.x0 + dx, bbox.y0 + dy, bbox.width, bbox.height])


@dataclass
class SetRcParams:
    """
    rcParams 的配置类型。继承此类，使用下划线命名法来填写 rcParams 的配置，
    然后使用 apply() 方法应用到 plt.rcParams 中。

    例如设置 rcParams['font.size'] = 10.5，可令 font_size = 10.5
    """
    def _underscore_to_dot(self, name: str) -> str:
        """驼峰命名转换为点分割命名"""
        return name.replace('_', '.')

    def apply(self):
        """将对象中的属性应用到plt.rcParams中
        """
        for key, value in asdict(self).items():
            key = self._underscore_to_dot(key)
            plt.rcParams[key] = value


@dataclass
class MoiRcParams(SetRcParams):
    """Moisten 的默认配置，适用于A4印刷"""
    font_size: float = 10.5

    # 边框线宽
    axes_linewidth: float = 0.8

    # 刻度线宽
    xtick_major_width: float = 0.8
    ytick_major_width: float = 0.8
    xtick_minor_width: float = 0.6
    ytick_minor_width: float = 0.6

    # 刻度线长度
    xtick_major_size: float = 4
    ytick_major_size: float = 4
    xtick_minor_size: float = 2.5
    ytick_minor_size: float = 2.5

    # 刻度线标签间隔
    xtick_major_pad: float = 3
    ytick_major_pad: float = 3

    # 刻度线标签大小
    xtick_labelsize: float = font_size
    ytick_labelsize: float = font_size


